import re
import top.utils as utils
import GROparser as gro
import PDBparser as pdb
import copy


class Topol:
    def __init__(self, itp_file='', mol='', res='', mode='aa'):
        self.topol_items = {
            'moleculetype': 'UNK',
            'atoms': [],
            'bonds': [],
            'pairs': [],
            'constraints': [],
            'angles': [],
            'dihedrals': [],
            'exclusions': []
        }
        self.mol = mol
        self.res = res
        self.mode = mode
        if itp_file != '':
            with open(itp_file) as f:
                lines = f.readlines()
                lines = [line.split(';')[0] for line in lines
                         if len(line.strip()) != 0 and
                         not line.strip().startswith(';') and
                         not line.strip().startswith('#')]
                for line in lines:
                    M = re.compile(r'.*\[(.*)].*')
                    if M.match(line):
                        cur_item = M.match(line).group(1).strip()
                    else:
                        if cur_item == 'moleculetype':
                            if mol != '':
                                self.topol_items[cur_item] = mol
                            else:
                                self.topol_items[cur_item] = line.split()[0]
                        if cur_item == 'atoms':
                            self.topol_items[cur_item].append(self.convert_atoms(line, res))
                        if cur_item == 'bonds':
                            self.topol_items[cur_item].append(self.convert_bonds(line))
                        if cur_item == 'pairs':
                            self.topol_items[cur_item].append(self.convert_pairs(line))
                        if cur_item == 'constraints':
                            self.topol_items[cur_item].append(self.convert_constraints(line))
                        if cur_item == 'angles':
                            self.topol_items[cur_item].append(self.convert_angles(line))
                        if cur_item == 'dihedrals':
                            self.topol_items[cur_item].append(self.convert_dihedrals(line))
                        if cur_item == 'exclusions':
                            self.topol_items[cur_item].append(self.convert_exclusions(line))
            print('Finished reading topology of', self.topol_items['moleculetype'] + '.')
        print(self)
        return

    def convert_atoms(self, line, res):
        if self.mode == 'aa':
            return TopolAtom(line, res)
        elif self.mode == 'cg':
            return TopolAtomCG(line, res)

    def convert_bonds(self, line):
        if self.mode == 'aa':
            return TopolBond(self.topol_items, line)
        elif self.mode == 'cg':
            return TopolBondCG(self.topol_items, line)

    def convert_pairs(self, line):
        return TopolPair(self.topol_items, line)

    def convert_constraints(self, line):
        return TopolConstraintCG(self.topol_items, line)

    def convert_angles(self, line):
        if self.mode == 'aa':
            return TopolAngle(self.topol_items, line)
        elif self.mode == 'cg':
            return TopolAngleCG(self.topol_items, line)

    def convert_dihedrals(self, line):
        if self.mode == 'aa':
            return TopolDihedral(self.topol_items, line)
        elif self.mode == 'cg':
            return TopolDihedralCG(self.topol_items, line)

    def convert_exclusions(self, line):
        if self.mode == 'aa':
            return TopolExclusion(self.topol_items, line)
        elif self.mode == 'cg':
            return TopolExclusionCG(self.topol_items, line)

    def __str__(self):
        return 'Molecule ' + self.topol_items['moleculetype'] + ' with \n' + \
               str(len(self.topol_items['atoms'])) + ' atoms \n' + \
               str(len(self.topol_items['bonds'])) + ' bonds \n' + \
               str(len(self.topol_items['pairs'])) + ' pairs \n' + \
               str(len(self.topol_items['angles'])) + ' angles \n' + \
               str(len(self.topol_items['dihedrals'])) + ' dihedrals \n' + \
               str(len(self.topol_items['exclusions'])) + ' exclusions \n'

    def shift_atoms(self, offset):
        # shift 原子编号
        for atom in self.topol_items['atoms']:
            if atom.nr > 0:
                atom.nr += offset
                atom.cgnr += offset
        print('Atom numbers of ' + self.topol_items['moleculetype'] + ' shifted by', offset)

    def shift_residues(self, offset):
        # shift 残基编号
        for atom in self.topol_items['atoms']:
            atom.resnr += offset
        print('Residue numbers of ' + self.topol_items['moleculetype'] + ' shifted by', offset)

    def mark_atom_list(self, mark_list_str, remove=False):
        self.__mark_atom__(utils.get_atom_list(mark_list_str), remove)

    def mark_atom_list_from_file(self, mark_list_file, remove=False):
        with open(mark_list_file) as f:
            self.mark_atom_list(f.readlines()[0], remove)

    def __mark_atom__(self, mark_list, remove=False):
        new_atom_list = []
        # 1. 添加保留的原子
        for atom in self.topol_items['atoms']:
            if atom.nr not in mark_list:
                new_atom_list.append(atom)
        # 2. 保留的原子重新排序
        for i, atom in enumerate(new_atom_list):
            atom.nr = i + 1
            atom.cgnr = i + 1
        # 3. 标记的原子的序号改为相反数, 添加到末尾
        for atom_id in mark_list:
            atom = get_atom_from_nr(self.topol_items, atom_id)
            atom.nr *= -1
            atom.cgnr *= -1
            new_atom_list.append(atom)
        self.topol_items['atoms'] = new_atom_list
        print('Marked atoms', mark_list, 'of', self.topol_items['moleculetype'] + '.')
        if remove:
            self.clean()

    def to_itp(self, file_name):
        print('Generating itp of ' + self.topol_items['moleculetype'] + '.')
        with open(file_name, 'w') as f:
            # 输出残基名
            f.writelines('[ moleculetype ]\n; Name   nrexcl\n')
            f.writelines(' ' + self.topol_items['moleculetype'] + '    3\n')
            # 输出原子编号
            f.writelines('[ atoms ]\n;  nr  type  resnr  resid  atom  cgnr  charge    mass\n')
            for atom in self.topol_items['atoms']:
                f.writelines(str(atom) + '\n')
            # 输出键长编号
            f.writelines('[ bonds ]\n;  ai   aj  funct   c0         c1\n')
            for bond in self.topol_items['bonds']:
                f.writelines(str(bond) + '\n')
            if self.mode == 'aa':
                # 输出 pair 编号
                f.writelines('[ pairs ]\n;  ai   aj  funct  ;  all 1-4 pairs but the ones excluded in GROMOS itp\n')
                for pair in self.topol_items['pairs']:
                    f.writelines(str(pair) + '\n')
            elif self.mode == 'cg':
                # 输出 constraint 编号
                f.writelines('[ constraints ]\n')
                for constraint in self.topol_items['constraints']:
                    f.writelines(str(constraint) + '\n')
            # 输出键角编号
            f.writelines('[ angles ]\n;  ai   aj   ak  funct   angle     fc\n')
            for angle in self.topol_items['angles']:
                f.writelines(str(angle) + '\n')
            # 输出二面角编号
            f.writelines('[ dihedrals ]\n; including dihedrals and impropers\n')
            for dihedral in self.topol_items['dihedrals']:
                f.writelines(str(dihedral) + '\n')
            # 输出 exclusion 编号
            f.writelines('[ exclusions ]\n;  ai   aj  funct  ;  GROMOS 1-4 exclusions\n')
            for exclusion in self.topol_items['exclusions']:
                f.writelines(str(exclusion) + '\n')

    def to_rtp(self, out_rtp, ff='gromos', exclude=''):
        contents = '; converted from top of ' + self.topol_items['moleculetype'] + '\n' + \
                   '; created by supernova/PaperTools (https://gitee.com/supernova_bingbing/paper-tools)\n\n'
        if ff == 'gromos':
            contents += '[ bondedtypes ]\n; bonds  angles  dihedrals  impropers\n' + \
                        '    2       2          1          2\n\n'
        elif ff == 'amber':
            contents += '[ bondedtypes ]\n' + \
                        '; bonds  angles  dihedrals  impropers all_dihedrals nrexcl HH14 RemoveDih\n' + \
                        '     1       1          9          4        1         3      1     0\n\n'
        else:
            contents = 'Fuck you!'

        # 排除项
        if exclude != '':
            with open(exclude) as f:
                lines = f.readlines()
                exclude_N = lines[0].strip()
                if exclude_N != '':
                    exclude_N = utils.get_atom_list(exclude_N)
                atom_N = lines[1].strip()
                if atom_N != '':
                    atom_N = int(atom_N)
                exclude_C = lines[2].strip()
                if exclude_C != '':
                    exclude_C = utils.get_atom_list(exclude_C)
                atom_C = lines[3].strip()
                if atom_C != '':
                    atom_C = int(atom_C)
        else:
            exclude_N = []
            atom_N = 0
            exclude_C = []
            atom_C = 0

        # 残基名
        contents += '[ {0:s} ]\n'.format(self.topol_items['moleculetype'])

        # [ atoms ]字段：记录残基中每个原子的名称、类型和电荷、电荷组
        contents += ' [ atoms ]\n'
        for atom in self.topol_items['atoms']:
            if exclude_N != '' and atom.nr in exclude_N:
                contents += ';- ' + atom.to_rtp_str() + '\n'
            elif exclude_C != '' and atom.nr in exclude_C:
                contents += ';+ ' + atom.to_rtp_str() + '\n'
            else:
                contents += atom.to_rtp_str() + '\n'

        # [ bonds ]字段：原子间的连接信息
        contents += ' [ bonds ]\n'
        for bond in self.topol_items['bonds']:
            if exclude_N != '' and (bond.ai.nr in exclude_N or bond.aj.nr in exclude_N):
                contents += ';- ' + bond.to_rtp_str() + '\n'
            elif exclude_C != '' and (bond.ai.nr in exclude_C or bond.aj.nr in exclude_C):
                contents += ';+ ' + bond.to_rtp_str() + '\n'
            else:
                contents += bond.to_rtp_str() + '\n'

        # [ angles ]字段：键角信息
        contents += ' [ angles ]\n'
        for angle in self.topol_items['angles']:
            if exclude_N != '' and (angle.ai.nr in exclude_N or angle.aj.nr in exclude_N or angle.ak.nr in exclude_N):
                contents += ';- ' + angle.to_rtp_str() + '\n'
            elif exclude_C != '' and (angle.ai.nr in exclude_C or angle.aj.nr in exclude_C or angle.ak.nr in exclude_C):
                contents += ';+ ' + angle.to_rtp_str() + '\n'
            else:
                contents += angle.to_rtp_str() + '\n'

        # [ dihedrals ]字段：二面角信息
        contents += ' [ dihedrals ]\n'
        for dihedral in self.topol_items['dihedrals']:
            if dihedral.type == 'dihedral':
                if exclude_N != '' and (dihedral.ai.nr in exclude_N or dihedral.aj.nr in exclude_N
                                        or dihedral.ak.nr in exclude_N or dihedral.al.nr in exclude_N):
                    contents += ';- ' + dihedral.to_rtp_str() + '\n'
                elif exclude_C != '' and (dihedral.ai.nr in exclude_C or dihedral.aj.nr in exclude_C
                                          or dihedral.ak.nr in exclude_C or dihedral.al.nr in exclude_C):
                    contents += ';+ ' + dihedral.to_rtp_str() + '\n'
                else:
                    contents += dihedral.to_rtp_str() + '\n'

        # [ impropers ]字段：异常二面角信息
        contents += ' [ impropers ]\n'
        for improper in self.topol_items['dihedrals']:
            if improper.type == 'improper':
                if exclude_N != '' and (improper.ai.nr in exclude_N or improper.aj.nr in exclude_N
                                        or improper.ak.nr in exclude_N or improper.al.nr in exclude_N):
                    contents += ';- ' + improper.to_rtp_str() + '\n'
                elif exclude_C != '' and (improper.ai.nr in exclude_C or improper.aj.nr in exclude_C
                                          or improper.ak.nr in exclude_C or improper.al.nr in exclude_C):
                    contents += ';+ ' + improper.to_rtp_str() + '\n'
                else:
                    contents += improper.to_rtp_str() + '\n'

        # [ exclusions ]字段：定义哪些原子之间的非键相互作用忽略掉
        contents += ' [ exclusions ]\n'
        for exclusion in self.topol_items['exclusions']:
            if exclude_N != '' and (exclusion.ai.nr in exclude_N or exclusion.aj.nr in exclude_N):
                contents += ';- ' + exclusion.to_rtp_str() + '\n'
            elif exclude_C != '' and (exclusion.ai.nr in exclude_C or exclusion.aj.nr in exclude_C):
                contents += ';+ ' + exclusion.to_rtp_str() + '\n'
            else:
                contents += exclusion.to_rtp_str() + '\n'

        with open(out_rtp, 'w') as f:
            f.writelines(contents)

        print('Generated rtp of ' + self.topol_items['moleculetype'] + '.')

    def append(self, next_itp, overlap=''):
        """ 合并两个分子的拓扑, 主体分子的itp实例调用客体分子的itp实例
            主体分子: ... A1 —— A1 —— A1 -- B1 —— B1 —— B1 ...
            客体分子: ... A2 —— A2 —— A2 -- B2 —— B2 —— B2 ...
            overlap: 文本文件, 包含3行, 定义主体分子的B1部分, 以及客体分子的B2和A2部分的原子编号 (A1不需要定义)
            格式:
                1-3, 5, 7-8
                4-7, 10, 12
                5-7
            其中第1行为主体B1部分原子, 第2行为客体B2部分原子, 第三行为客体A2部分原子
            NOTE: 第一行和第二行的原子顺序、对应关系一定要严格一致!!! """
        if overlap != '':
            with open(overlap) as f:
                lines = f.readlines()
                over_b1 = utils.get_atom_list(lines[0])
                over_b2 = utils.get_atom_list(lines[1])
                over_a2 = utils.get_atom_list(lines[2])
                if len(over_b1) != len(over_b2):
                    raise Exception("Group B1 and B2 must be of one-to-one correspondence.")
        else:
            over_b1 = []
            over_b2 = []
            over_a2 = []

        # 使用深拷贝建立next_itp的实例, 可以多次调用本方法
        next_itp = copy.deepcopy(next_itp)

        # Combine bonds part
        print('Appending bonds of', next_itp.topol_items['moleculetype'], 'into',
              self.topol_items['moleculetype'] + '...')
        new_bonds_list = []
        # 保留主体所有包含B1的项
        for bond in self.topol_items['bonds']:
            if len(over_b1) != 0 and len(over_b2) != 0:
                # 排除 B1-B1
                if bond.ai.nr not in over_b1 or bond.aj.nr not in over_b1:
                    if bond.ai.nr in over_b1:
                        bond.ai = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(bond.ai.nr)])
                    if bond.aj.nr in over_b1:
                        bond.aj = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(bond.aj.nr)])
                    new_bonds_list.append(bond)
            else:
                new_bonds_list.append(bond)
        # 删除客体所有包含A2的项
        for bond in next_itp.topol_items['bonds']:
            if len(over_a2) != 0:
                # 排除 A2-A2
                if bond.ai.nr not in over_a2 and bond.aj.nr not in over_a2:
                    new_bonds_list.append(bond)
            else:
                new_bonds_list.append(bond)
        self.topol_items['bonds'] = new_bonds_list

        # Combine pairs part
        print('Appending pairs of', next_itp.topol_items['moleculetype'], 'into',
              self.topol_items['moleculetype'] + '...')
        new_pairs_list = []
        # 保留主体所有包含B1的项
        for pair in self.topol_items['pairs']:
            if len(over_b1) != 0 and len(over_b2) != 0:
                # B1-B1
                if pair.ai.nr not in over_b1 or pair.aj.nr not in over_b1:
                    if pair.ai.nr in over_b1:
                        pair.ai = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(pair.ai.nr)])
                    if pair.aj.nr in over_b1:
                        pair.aj = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(pair.aj.nr)])
                    new_pairs_list.append(pair)
            else:
                new_pairs_list.append(pair)
        # 删除客体所有包含A2的项
        for pair in next_itp.topol_items['pairs']:
            if len(over_a2) != 0:
                # A2-A2
                if pair.ai.nr not in over_a2 and pair.aj.nr not in over_a2:
                    new_pairs_list.append(pair)
            else:
                new_pairs_list.append(pair)
        self.topol_items['pairs'] = new_pairs_list

        # Combine angles part
        print('Appending angles of', next_itp.topol_items['moleculetype'], 'into',
              self.topol_items['moleculetype'] + '...')
        new_angles_list = []
        # 保留主体所有包含B1的项
        for angle in self.topol_items['angles']:
            if len(over_b1) != 0 and len(over_b2) != 0:
                # B1-B1-B1
                if angle.ai.nr not in over_b1 or angle.aj.nr not in over_b1 or angle.ak.nr not in over_b1:
                    if angle.ai.nr in over_b1:
                        angle.ai = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(angle.ai.nr)])
                    if angle.aj.nr in over_b1:
                        angle.aj = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(angle.aj.nr)])
                    if angle.ak.nr in over_b1:
                        angle.ak = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(angle.ak.nr)])
                    new_angles_list.append(angle)
            else:
                new_angles_list.append(angle)
        # 删除客体所有包含A2的项
        for angle in next_itp.topol_items['angles']:
            if len(over_a2) != 0:
                # A2-A2-A2
                if angle.ai.nr not in over_a2 and angle.aj.nr not in over_a2 and angle.ak.nr not in over_a2:
                    new_angles_list.append(angle)
            else:
                new_angles_list.append(angle)
        self.topol_items['angles'] = new_angles_list

        # Combine dihedrals part
        print('Appending dihedrals of', next_itp.topol_items['moleculetype'], 'into',
              self.topol_items['moleculetype'] + '...')
        new_dihedrals_list = []
        # 保留主体所有包含B1的项
        for dihedral in self.topol_items['dihedrals']:
            if len(over_b1) != 0 and len(over_b2) != 0:
                # B1-B1-B1-B1
                if dihedral.ai.nr not in over_b1 or dihedral.aj.nr not in over_b1 \
                        or dihedral.ak.nr not in over_b1 or dihedral.al.nr not in over_b1:
                    if dihedral.ai.nr in over_b1:
                        dihedral.ai = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(dihedral.ai.nr)])
                    if dihedral.aj.nr in over_b1:
                        dihedral.aj = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(dihedral.aj.nr)])
                    if dihedral.ak.nr in over_b1:
                        dihedral.ak = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(dihedral.ak.nr)])
                    if dihedral.al.nr in over_b1:
                        dihedral.al = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(dihedral.al.nr)])
                    new_dihedrals_list.append(dihedral)
            else:
                new_dihedrals_list.append(dihedral)
        # 删除客体所有包含A2的项
        for dihedral in next_itp.topol_items['dihedrals']:
            if len(over_a2) != 0:
                # A2-A2-A2-A2
                if dihedral.ai.nr not in over_a2 and dihedral.aj.nr not in over_a2 \
                        and dihedral.ak.nr not in over_a2 and dihedral.al.nr not in over_a2:
                    new_dihedrals_list.append(dihedral)
            else:
                new_dihedrals_list.append(dihedral)
        self.topol_items['dihedrals'] = new_dihedrals_list

        # Combine exclusions part
        print('Appending exclusions of', next_itp.topol_items['moleculetype'], 'into',
              self.topol_items['moleculetype'] + '...')
        new_exclusions_list = []
        # 保留主体所有包含B1的项
        for exclusion in self.topol_items['exclusions']:
            if len(over_b1) != 0 and len(over_b2) != 0:
                # B1-B1
                if exclusion.ai.nr not in over_b1 or exclusion.aj.nr not in over_b1:
                    if exclusion.ai.nr in over_b1:
                        exclusion.ai = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(exclusion.ai.nr)])
                    if exclusion.aj.nr in over_b1:
                        exclusion.aj = get_atom_from_nr(next_itp.topol_items, over_b2[over_b1.index(exclusion.aj.nr)])
                    new_exclusions_list.append(exclusion)
            else:
                new_exclusions_list.append(exclusion)
        # 删除客体所有包含A2的项
        for exclusion in self.topol_items['exclusions']:
            if len(over_a2) != 0:
                # A2-A2
                if exclusion.ai.nr not in over_a2 and exclusion.aj.nr not in over_a2:
                    new_exclusions_list.append(exclusion)
            else:
                new_exclusions_list.append(exclusion)
        self.topol_items['exclusions'] = new_exclusions_list

        # Combine atoms part
        if len(over_b1) > 0:
            next_itp.shift_residues(get_atom_from_nr(self.topol_items, over_b1[0] - 1).resnr - 1)  # 客体残基编号平移
        else:
            next_itp.shift_residues(self.topol_items['atoms'][-1].resnr)  # 客体残基编号平移
        print('Appending atoms of', next_itp.topol_items['moleculetype'], 'into',
              self.topol_items['moleculetype'] + '...')
        new_atoms_list = []
        for pair in self.topol_items['atoms']:
            if pair.nr not in over_b1:
                new_atoms_list.append(pair)
        for pair in next_itp.topol_items['atoms']:
            if pair.nr not in over_a2:
                new_atoms_list.append(pair)
        self.topol_items['atoms'] = new_atoms_list

        self.__resort__()

    def add_atom(self, nr, typ, resnr, resid, atom, cgnr, charge, mass, shift_forward_residues=False):
        atom = TopolAtom('{0:5d}{1:>6s}{2:5d}{3:>7s}{4:>8s}{5:5d}{6:12.6f}{7:9.4f}'.format(nr, typ, resnr, resid, atom, cgnr, charge, mass), resid)
        if shift_forward_residues:
            for nr in range(atom.nr, len(self.topol_items['atoms']) + 1):
                get_atom_from_nr(self.topol_items, nr).resnr += 1
        self.topol_items['atoms'].insert(atom.nr - 1, atom)
        self.__resort__()

    def add_bond(self, atnr1, atnr2, funct, c0=0, c1=0):
        if c0 == 0 and c1 == 0:
            bond = TopolBond(self.topol_items, '{0:5d}{1:5d}{2:5d}'.format(atnr1, atnr2, funct))
        else:
            bond = TopolBond(self.topol_items, '{0:5d}{1:5d}{2:5d}{3:9.4f}{4:13.4e}'.format(atnr1, atnr2, funct, c0, c1))
        self.topol_items['bonds'].append(bond)

    def add_pair(self, atnr1, atnr2, funct):
        pair = TopolPair(self.topol_items, '{0:5d}{1:5d}{2:5d}'.format(atnr1, atnr2, funct))
        self.topol_items['pairs'].append(pair)

    def add_angle(self, atnr1, atnr2, atnr3, funct, c0=0, c1=0):
        if c0 == 0 and c1 == 0:
            angle = TopolAngle(self.topol_items, '{0:5d}{1:5d}{2:5d}{3:5d}' \
                               .format(atnr1, atnr2, atnr3, funct))
        else:
            angle = TopolAngle(self.topol_items, '{0:5d}{1:5d}{2:5d}{3:5d}{4:10.2f}{5:9.2f}' \
                    .format(atnr1, atnr2, atnr3, funct, c0, c1))
        self.topol_items['angles'].append(angle)

    def add_dihedral(self, atnr1, atnr2, atnr3, atnr4, funct, c0, c1, c2=0):
        if funct == 2:
            line = '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}' \
                .format(atnr1, atnr2, atnr3, atnr4, funct, c0, c1)
        elif funct == 4:
            line = '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                .format(atnr1, atnr2, atnr3, atnr4, funct, c0, c1, c2)
        elif funct == 1:
            line = '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                .format(atnr1, atnr2, atnr3, atnr4, funct, c0, c1, c2)
        elif funct == 9:
            line = '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                .format(atnr1, atnr2, atnr3, atnr4, funct, c0, c1, c2)
        else:
            raise Exception('Unknown dihedral type.')
        dihedral = TopolDihedral(self.topol_items, line)
        self.topol_items['dihedrals'].append(dihedral)

    def add_exclusion(self, atnr1, atnr2):
        exclusion = TopolExclusion(self.topol_items, '{0:5d}{1:5d}'.format(atnr1, atnr2))
        self.topol_items['exclusions'].append(exclusion)

    def sort_atoms_by_pdb(self, ref_pdb):
        if not ref_pdb.endswith('.pdb'):
            raise Exception('Input reference is not in .pdb format.')
        # Waiting for pdb reader
        ref = pdb.PDB(ref_pdb)
        l = [[atom.resid, atom.__atid_for_human__(atom.atid)] for atom in ref.atoms]
        self.__sort_atoms_by_ref__(l)

    def sort_atoms_by_gro(self, ref_gro):
        if not ref_gro.endswith('.gro'):
            raise Exception('Input reference is not in .gro format.')
        ref = gro.GRO(ref_gro)
        l = [[atom.resid, atom.atname] for atom in ref.atom_list]
        self.__sort_atoms_by_ref__(l)

    def __sort_atoms_by_ref__(self, res_atom_list):
        """拓扑原子重新排序
        paras: res_atom_list list[残基编号. 原子名]
        return: None
        """
        print('Sorting', str(len(self.topol_items['atoms'])), 'atoms.')
        self.topol_items['atoms'] = [get_nr_from_atom(self.topol_items, res, atname) for (res, atname) in res_atom_list]
        for i, atom in enumerate(self.topol_items['atoms']):
            atom.nr = i + 1

    def __resort__(self):
        """强制重新排序"""
        for i, atom in enumerate(self.topol_items['atoms']):
            atom.nr = i + 1
            atom.cgnr = i + 1

    def clean(self):
        print('Cleaning useless atoms of', self.topol_items['moleculetype'] + '.')

        new_atom_list = []
        for atom in self.topol_items['atoms']:
            if atom.nr > 0:
                new_atom_list.append(atom)
        self.topol_items['atoms'] = new_atom_list

        new_bond_list = []
        for bond in self.topol_items['bonds']:
            if bond.ai.nr > 0 and bond.aj.nr > 0:
                new_bond_list.append(bond)
        self.topol_items['bonds'] = new_bond_list

        new_pair_list = []
        for pair in self.topol_items['pairs']:
            if pair.ai.nr > 0 and pair.aj.nr > 0:
                new_pair_list.append(pair)
        self.topol_items['pairs'] = new_pair_list

        new_angle_list = []
        for angle in self.topol_items['angles']:
            if angle.ai.nr > 0 and angle.aj.nr > 0 and angle.ak.nr > 0:
                new_angle_list.append(angle)
        self.topol_items['angles'] = new_angle_list

        new_dihedral_list = []
        for dihedral in self.topol_items['dihedrals']:
            if dihedral.ai.nr > 0 and dihedral.aj.nr > 0 \
                    and dihedral.ak.nr > 0 and dihedral.al.nr > 0:
                new_dihedral_list.append(dihedral)
        self.topol_items['dihedrals'] = new_dihedral_list

        new_exclusion_list = []
        for exclusion in self.topol_items['exclusions']:
            if exclusion.ai.nr > 0 and exclusion.aj.nr > 0:
                new_exclusion_list.append(exclusion)
        self.topol_items['exclusions'] = new_exclusion_list


def get_atom_from_nr(topol_items, nr):
    """从原子编号获取原子拓扑对象"""
    for atom in topol_items['atoms']:
        if atom.nr == nr:
            return atom


def get_nr_from_atom(topol_items, resnr, atname):
    """从原子拓扑对象获取原子编号"""
    for atom in topol_items['atoms']:
        if atom.resnr == resnr and atom.atom == atname:
            return atom


class TopolAtom:
    # nr  type  resnr  resid  atom  cgnr  charge    mass
    def __init__(self, line, res=''):
        nr, typ, resnr, resid, atom, cgnr, charge, mass = line.split()
        self.nr = int(nr)
        self.type = typ
        self.resnr = int(resnr)
        self.resid = res or resid
        self.atom = atom
        self.cgnr = int(cgnr)
        self.charge = float(charge)
        self.mass = float(mass)

    def to_rtp_str(self):
        # atom  type   charge  cgnr
        return '{0:>8s}{1:>6s}{2:9.3f}{3:5d}' \
            .format(self.atom, self.type, self.charge, self.cgnr)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1  HS14    1    FMC     HN2    1    0.245   1.0080
        return '{0:5d}{1:>6s}{2:5d}{3:>7s}{4:>8s}{5:5d}{6:12.6f}{7:9.4f}' \
            .format(self.nr, self.type, self.resnr, self.resid, self.atom, self.cgnr, self.charge, self.mass)


class TopolAtomCG:
    # nr  type  resnr  resid  atom  cgnr  charge
    def __init__(self, line, res=''):
        nr, type, resnr, resid, atom, cgnr, charge = line.split()
        self.nr = int(nr)
        self.type = type
        self.resnr = int(resnr)
        if res != '':
            self.resid = res
        else:
            self.resid = resid
        self.atom = atom
        self.cgnr = int(cgnr)
        self.charge = float(charge)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    P5     1   LYS    BB     1  0.0000 ; C
        return '{0:5d}{1:>6s}{2:5d}{3:>7s}{4:>8s}{5:5d}{6:9.3f}' \
            .format(self.nr, self.type, self.resnr, self.resid, self.atom, self.cgnr, self.charge)


class TopolBond:
    #   ai   aj  funct   c0         c1
    def __init__(self, topol_items, line):
        self.topol_items = topol_items
        paras = line.split()
        ai, aj, funct = paras[:3]
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.funct = int(funct)
        if len(paras) > 3:
            c0, c1 = paras[3:]
            self.c0 = float(c0)
            self.c1 = float(c1)
        else:
            self.c0 = None
            self.c1 = None

    def to_rtp_str(self):
        # ai   aj   c0         c1
        if self.c0 is not None and self.c1 is not None:
            return '{0:>5s}{1:>5s}{2:9.4f}{3:13.4e}' \
                .format(self.ai.atom, self.aj.atom, self.c0, self.c1)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    2    2   0.0985   8.2455e+07
        if self.c0 is not None and self.c1 is not None:
            return '{0:5d}{1:5d}{2:5d}{3:9.4f}{4:13.4e}' \
                .format(self.ai.nr, self.aj.nr, self.funct, self.c0, self.c1)
        else:
            return '{0:5d}{1:5d}{2:5d}' \
                .format(self.ai.nr, self.aj.nr, self.funct)


class TopolBondCG:
    #   ai   aj  funct   c0         c1
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        ai, aj, funct, c0, c1 = str.split()
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.funct = int(funct)
        self.c0 = float(c0)
        self.c1 = float(c1)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    2    2   0.0985   8.2455e+07
        return '{0:5d}{1:5d}{2:5d}{3:10.5f}{4:13.4e}' \
            .format(self.ai.nr, self.aj.nr, self.funct, self.c0, self.c1)


class TopolPair:
    #   ai   aj  funct
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        ai, aj, funct = str.split()
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.funct = int(funct)

    def to_rtp_str(self):
        #   ai   aj
        return '{0:>5s}{1:>5s}' \
            .format(self.ai.atom, self.aj.atom)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    5    1
        return '{0:5d}{1:5d}{2:5d}' \
            .format(self.ai.nr, self.aj.nr, self.funct)


class TopolConstraintCG:
    #   ai   aj  funct c0
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        ai, aj, funct, c0 = str.split()
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.funct = int(funct)
        self.c0 = float(c0)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    5    1
        return '{0:5d}{1:5d}{2:5d}{3:10.5f}' \
            .format(self.ai.nr, self.aj.nr, self.funct, self.c0)


class TopolAngle:
    #  ai   aj   ak  funct   angle     fc
    def __init__(self, topol_items, line):
        self.topol_items = topol_items
        paras = line.split()
        ai, aj, ak, funct = paras[:4]
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.ak = get_atom_from_nr(topol_items, int(ak))
        self.funct = int(funct)
        if len(paras) > 4:
            angle, fc = paras[4:]
            self.angle = float(angle)
            self.fc = float(fc)
        else:
            self.angle = None
            self.fc = None

    def to_rtp_str(self):
        if self.angle is not None and self.fc is not None:
            return '{0:>5s}{1:>5s}{2:>5s}{3:10.2f}{4:9.2f}' \
                .format(self.ai.atom, self.aj.atom, self.ak.atom, self.angle, self.fc)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    2    3    2    120.00   445.00
        if self.angle is not None and self.fc is not None:
            return '{0:5d}{1:5d}{2:5d}{3:5d}{4:10.2f}{5:9.2f}' \
                .format(self.ai.nr, self.aj.nr, self.ak.nr, self.funct, self.angle, self.fc)
        else:
            return '{0:5d}{1:5d}{2:5d}{3:5d}' \
                .format(self.ai.nr, self.aj.nr, self.ak.nr, self.funct)


class TopolAngleCG:
    #  ai   aj   ak  funct   angle     fc
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        ai, aj, ak, funct, angle, fc = str.split()
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.ak = get_atom_from_nr(topol_items, int(ak))
        self.funct = int(funct)
        self.angle = float(angle)
        self.fc = float(fc)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    2    3    2    120.00   445.00
        return '{0:5d}{1:5d}{2:5d}{3:5d}{4:10.2f}{5:9.2f}' \
            .format(self.ai.nr, self.aj.nr, self.ak.nr, self.funct, self.angle, self.fc)


class TopolDihedral:
    # improper  ai   aj   ak   al  funct   angle     fc
    # dihedral  ai   aj   ak   al  funct    ph0      cp     mult
    def __init__(self, topol_items, line):
        paras = line.split()
        ai, aj, ak, al, funct = paras[:5]
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))
        self.ak = get_atom_from_nr(topol_items, int(ak))
        self.al = get_atom_from_nr(topol_items, int(al))
        self.funct = int(funct)
        # 2和4是improper, 2有两个参数angle, fc, 4有3个参数angle, fc, mult
        if self.funct == 2:
            if len(paras) > 5:
                angle, fc = paras[5:]
                self.angle = float(angle)
                self.fc = float(fc)
                self.type = 'improper'
            else:
                self.angle = None
                self.fc = None
                self.type = 'improper'
        elif self.funct == 4:
            if len(paras) > 5:
                angle, fc, mult = paras[5:]
                self.angle = float(angle)
                self.fc = float(fc)
                self.mult = int(mult)
                self.type = 'improper'
            else:
                self.angle = None
                self.fc = None
                self.mult = None
                self.type = 'improper'
        elif self.funct == 1:
            if len(paras) > 5:
                ph0, cp, mult = paras[5:]
                self.ph0 = float(ph0)
                self.cp = float(cp)
                self.mult = int(mult)
                self.type = 'dihedral'
            else:
                self.ph0 = None
                self.cp = None
                self.mult = None
                self.type = 'dihedral'
        elif self.funct == 9:
            if len(paras) == 6:
                self.appendix = paras[5]
                self.phase = None
                self.kd = None
                self.pn = None
            elif len(paras) > 6:
                phase, kd, pn = paras[5:]
                self.phase = float(phase)
                self.kd = float(kd)
                self.pn = int(pn)
                self.type = 'dihedral'
            else:
                self.phase = None
                self.kd = None
                self.pn = None
                self.type = 'dihedral'
        else:
            raise Exception('Fuck')

    def to_rtp_str(self):
        if self.funct == 2:
            return '{0:>5s}{1:>5s}{2:>5s}{3:>5s}{4:10.2f}{5:9.2f}' \
                .format(self.ai.atom, self.aj.atom, self.ak.atom, self.al.atom, self.angle, self.fc)
        elif self.funct == 4:
            return '{0:>5s}{1:>5s}{2:>5s}{3:>5s}{4:10.2f}{5:9.2f}{6:9d}' \
                .format(self.ai.atom, self.aj.atom, self.ak.atom, self.al.atom, self.angle, self.fc, self.mult)
        elif self.funct == 1:
            return '{0:>5s}{1:>5s}{2:>5s}{3:>5s}{4:10.2f}{5:9.2f}{6:9d}' \
                .format(self.ai.atom, self.aj.atom, self.ak.atom, self.al.atom, self.ph0, self.cp, self.mult)
        elif self.funct == 9:
            return '{0:>5s}{1:>5s}{2:>5s}{3:>5s}{4:10.2f}{5:9.2f}{6:9d}' \
                .format(self.ai.atom, self.aj.atom, self.ak.atom, self.al.atom, self.phase, self.kd, self.pn)
        else:
            raise Exception('Unknown dihedral type.')

    def __str__(self):
        if self.funct == 2:
            if self.angle is not None and self.fc is not None:
                # 12345678901234567890123456789012345678901234567890123456789
                #     9    8   10   17    2      0.00   167.36
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct, self.angle, self.fc)
            else:
                # 12345678901234567890123456789012345678901234567890123456789
                #     9    8   10   17    2      0.00   167.36
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct)
        elif self.funct == 4:
            if self.angle is not None and self.fc is not None and self.mult is not None:
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct, self.angle, self.fc, self.mult)
            else:
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct)
        elif self.funct == 1:
            if self.ph0 is not None and self.cp is not None and self.mult is not None:
                # 12345678901234567890123456789012345678901234567890123456789
                #     3    2    4    6    1    180.00    33.50    2
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct, self.ph0, self.cp, self.mult)
            else:
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct)
        elif self.funct == 9:
            if self.phase is not None and self.kd is not None and self.pn is not None:
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct, self.phase, self.kd, self.pn)
            else:
                return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}' \
                    .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct)
        else:
            raise Exception('Unknown dihedral type.')


class TopolDihedralCG:
    # improper  ai   aj   ak   al  funct   angle     fc
    # dihedral  ai   aj   ak   al  funct    ph0      cp     mult
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        if len(str.split()) == 7:
            ai, aj, ak, al, funct, angle, fc = str.split()
            self.ai = get_atom_from_nr(topol_items, int(ai))
            self.aj = get_atom_from_nr(topol_items, int(aj))
            self.ak = get_atom_from_nr(topol_items, int(ak))
            self.al = get_atom_from_nr(topol_items, int(al))
            self.funct = int(funct)
            self.angle = float(angle)
            self.fc = float(fc)
            self.type = 'improper'
        else:
            ai, aj, ak, al, funct, ph0, cp, mult = str.split()
            self.ai = get_atom_from_nr(topol_items, int(ai))
            self.aj = get_atom_from_nr(topol_items, int(aj))
            self.ak = get_atom_from_nr(topol_items, int(ak))
            self.al = get_atom_from_nr(topol_items, int(al))
            self.funct = int(funct)
            self.ph0 = float(ph0)
            self.cp = float(cp)
            self.mult = int(mult)
            self.type = 'dihedral'

    def __str__(self):
        if self.type == 'improper':
            # 12345678901234567890123456789012345678901234567890123456789
            #     9    8   10   17    2      0.00   167.36
            return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}' \
                .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct, self.angle, self.fc)
        else:
            # 12345678901234567890123456789012345678901234567890123456789
            #     3    2    4    6    1    180.00    33.50    2
            return '{0:5d}{1:5d}{2:5d}{3:5d}{4:5d}{5:10.2f}{6:9.2f}{7:9d}' \
                .format(self.ai.nr, self.aj.nr, self.ak.nr, self.al.nr, self.funct, self.ph0, self.cp, self.mult)


class TopolExclusion:
    #   ai   aj
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        ai, aj = str.split()
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))

    def to_rtp_str(self):
        return '{0:>5s}{1:>5s}' \
            .format(self.ai.atom, self.aj.atom)

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    5
        return '{0:5d}{1:5d}' \
            .format(self.ai.nr, self.aj.nr)


class TopolExclusionCG:
    #   ai   aj
    def __init__(self, topol_items, str):
        self.topol_items = topol_items
        ai, aj = str.split()
        self.ai = get_atom_from_nr(topol_items, int(ai))
        self.aj = get_atom_from_nr(topol_items, int(aj))

    def __str__(self):
        # 12345678901234567890123456789012345678901234567890123456789
        #     1    5
        return '{0:5d}{1:5d}' \
            .format(self.ai.nr, self.aj.nr)


if __name__ == '__main__':
    top_left = Topol('../test/top/test_left.itp')
    top_right = Topol('../test/top/test_right.itp')
    top_left.append(top_right, overlap='../test/top/overlap.txt')
    top_left.to_itp('../test/top/test_total.itp')
